# -*- coding: utf-8 -*-

import os
import math
import random
import h5py
import cv2
import numpy as np
import scipy.io


def show_ground_truth(ground_truth, name="color_map"):
    min_value, max_value = np.min(ground_truth), np.max(ground_truth)
    img = ((ground_truth - min_value) / max_value) * 255
    img = img.astype(np.uint8, copy=True)
    img = cv2.resize(img, (img.shape[1]*4, img.shape[0]*4))
    color_map = cv2.applyColorMap(img, cv2.COLORMAP_JET)
    cv2.imshow(name, color_map)
    cv2.waitKey()


def point_distance(a, b):
    assert len(a) == len(b) and len(a) == 2
    return math.sqrt((a[0] - b[0])**2 + (a[1] - b[1])**2)


def get_crop_shape(shape):
    assert len(shape) == 2
    h, w = shape[0] // 4, shape[1] // 4
    h = int(math.ceil(h/4) * 4)
    w = int(math.ceil(w/4) * 4)
    assert h % 4 == 0 and w % 4 == 0
    return [h, w]


def generate_ground_truth(mat_path, shape, image_path=None):
    shape = [shape[0]//4, shape[1]//4]
    mat = scipy.io.loadmat(mat_path)
    locations = mat["image_info"][0][0][0][0][0]
    locations = [[location[0]/8, location[1]/8] for location in locations]
    ground_truth = np.zeros(shape=shape, dtype=np.float32)
    for location in locations:
        temp = np.zeros_like(ground_truth)
        x, y = math.floor(location[0]), math.floor(location[1])
        if x >= shape[1] or y >= shape[0]:
            continue
        temp[y, x] = 1
        temp = cv2.GaussianBlur(temp, (0, 0), 5)
        ground_truth += temp
    # print(len(locations), np.sum(ground_truth))
    if image_path:
        print(len(locations), np.sum(ground_truth))
        image = cv2.imread(image_path)
        image = cv2.resize(image, dsize=(shape[1], shape[0]))
        cv2.imshow("src", image)
        show_ground_truth(ground_truth)
        cv2.destroyAllWindows()
    return ground_truth

"""
def main():
    folder = "/home/lijun/data/dataset/DataSet/ShanghaiTech/part_B_final/train_data/images"
    images_name = [name for name in os.listdir(folder) if "jpg" in name]
    save_folder = "/home/lijun/Dataset/crowd_counting/part_b/train_multi_10"
    save_index = 1
    for name in images_name:
        image_path = os.path.join(folder, name)
        mat_path = image_path.replace("images/", "ground_truth/GT_").replace("jpg", "mat")
        image = cv2.imread(image_path)
        image = cv2.resize(image, dsize=(512, 384))
        img_height, img_width = image.shape[0], image.shape[1]
        ground_truth = generate_ground_truth(mat_path, [img_height, img_width], image_path=None)
        ground_truth *= 10
        ground_truth_height = ground_truth.shape[0]
        ground_truth_width = ground_truth.shape[1]
        assert ground_truth_height == 96 and ground_truth_width == 128
        assert img_height == 384 and img_width == 512
        img_crop_height, img_crop_width = img_height // 2, img_width // 2
        ground_truth_crop_height, ground_truth_crop_width = ground_truth_height // 2, ground_truth_width // 2
        for y in range(0, 2):
            img_offset_y = y * (img_height // 2)
            ground_truth_offset_y = y * (ground_truth_height // 2)
            for x in range(0, 2):
                img_offset_x = x * (img_width // 2)
                ground_truth_offset_x = x * (ground_truth_width // 2)
                img_crop = image[img_offset_y: img_offset_y + img_crop_height,
                                 img_offset_x: img_offset_x + img_crop_width, :]
                ground_truth_crop = ground_truth[
                                    ground_truth_offset_y: ground_truth_offset_y + ground_truth_crop_height,
                                    ground_truth_offset_x: ground_truth_offset_x + ground_truth_crop_width]
                img_save_path = "{}/images/{:>07}.jpg".format(save_folder, save_index)
                ground_truth_save_path = "{}/ground_truth/{:>07}.h5".format(save_folder, save_index)
                cv2.imwrite(img_save_path, img_crop)
                h5_file = h5py.File(ground_truth_save_path, "w")
                h5_file.create_dataset("data", data=ground_truth_crop)
                h5_file.close()
                save_index += 1
                img_save_path = "{}/images/{:>07}.jpg".format(save_folder, save_index)
                ground_truth_save_path = "{}/ground_truth/{:>07}.h5".format(save_folder, save_index)
                img_crop = cv2.flip(img_crop, flipCode=1)  # horizontal flip
                ground_truth_crop = cv2.flip(ground_truth_crop, flipCode=1)  # horizontal flip
                cv2.imwrite(img_save_path, img_crop)
                h5_file = h5py.File(ground_truth_save_path, "w")
                h5_file.create_dataset("data", data=ground_truth_crop)
                h5_file.close()
                save_index += 1
        print(save_index // 8)
"""

"""
def main():
    folder = "/home/lijun/data/dataset/DataSet/ShanghaiTech/part_B_final/test_data/images"
    images_name = [name for name in os.listdir(folder) if "jpg" in name]
    min_values, max_values = [], []
    for name in images_name:
        image_path = os.path.join(folder, name)
        mat_path = image_path.replace("images/", "ground_truth/GT_").replace("jpg", "mat")
        image = cv2.imread(image_path)
        image = cv2.resize(image, dsize=(512, 384))
        # cv2.imshow("show", image)
        img_height, img_width = image.shape[0], image.shape[1]
        ground_truth = generate_ground_truth(mat_path, [img_height, img_width], image_path=None)
        ground_truth *= 10
        min_values.append(np.min(ground_truth))
        max_values.append(np.max(ground_truth))
        # show_ground_truth(ground_truth)
    cv2.destroyAllWindows()
    min_values = sorted(min_values)
    max_values = sorted(max_values)
    print("min value: ", min_values[0], min_values[-1])
    print("max value: ", max_values[0], max_values[-1])
"""


def main():
    folder = "/home/lijun/Dataset/crowd_counting/part_b/train_multi_1000/images"
    images_name = [name for name in os.listdir(folder) if "jpg" in name]
    for _ in range(1, 20):
        name = random.choice(images_name)
        image_path = os.path.join(folder, name)
        mat_path = image_path.replace("images", "ground_truth").replace("jpg", "h5")
        image = cv2.imread(image_path)
        ground_truth_file = h5py.File(mat_path, "r")
        ground_truth = np.array(ground_truth_file["data"])
        print(np.sum(ground_truth))
        ground_truth_file.close()
        cv2.imshow("image", image)
        print(np.min(ground_truth), np.max(ground_truth))
        show_ground_truth(ground_truth)
    cv2.destroyAllWindows()


"""
def main():
    image_path = "/home/lijun/data/dataset/DataSet/ShanghaiTech/part_B_final/train_data/images/IMG_4.jpg"
    image = cv2.imread(image_path)
    width, height = image.shape[1] // 4, image.shape[0] // 4
    image = cv2.resize(image, dsize=(width, height))
    mat_path = image_path.replace("images/", "ground_truth/GT_").replace("jpg", "mat")
    mat = scipy.io.loadmat(mat_path)
    locations = mat["image_info"][0][0][0][0][0]
    locations = [[location[0]/4, location[1]/4] for location in locations]
    for location in locations:
        cv2.circle(image, (int(location[0]), int(location[1])), 2, (0, 0, 255), -1)
    cv2.imshow("show", image)
    cv2.waitKey()
"""


if __name__ == "__main__":
    main()
